#!/usr/bin/env python

from SocketServer import *
from socket import *
import sys, os
from struct import *
import time
import threading
import thread

def daemonize (stdin='/dev/null', stdout='/dev/null', stderr='/dev/null'): 
    # Perform first fork. 
    try: 
        pid = os.fork( ) 
        if pid > 0: 
            sys.exit(0) # Exit first parent. 
    except OSError, e: 
        sys.stderr.write("fork #1 failed: (%d) %sn" % (e.errno, e.strerror)) 
        sys.exit(1) 
            
    # Decouple from parent environment. 
    os.chdir("/") 
    os.umask(0) 
    os.setsid( ) 
        
    # Perform second fork. 
    try: 
        pid = os.fork( ) 
        if pid > 0: 
            sys.exit(0) # Exit second parent. 
    except OSError, e: 
        sys.stderr.write("fork #2 failed: (%d) %sn" % (e.errno, e.strerror)) 
        sys.exit(1) 
    
    # The process is now daemonized, redirect standard file descriptors. 
    for f in sys.stdout, sys.stderr: f.flush( ) 
    si = file(stdin, 'r') 
    so = file(stdout, 'a+') 
    se = file(stderr, 'a+', 0) 
    os.dup2(si.fileno( ), sys.stdin.fileno( )) 
    os.dup2(so.fileno( ), sys.stdout.fileno( )) 
    os.dup2(se.fileno( ), sys.stderr.fileno( ))
                                                                                                                                                                
DEF_LOCAL_HOST = '0.0.0.0'
#DEF_LOCAL_HOST = '192.168.3.1'
DEF_REMOTE_SERVER = '8.8.8.8'
DEF_REMOTE_PORT = 53
DEF_LOCAL_PORT = 1054
#DEF_TIMEOUT = 0.4
DEF_TIMEOUT = 1

gl_remote_server = None

gl_mutex = threading.Lock()
gl_sock_locks = []

def thread_get_sock(socks, sock_index):  

    if ( gl_sock_locks[sock_index].acquire(False)):
        gl_sock_locks[sock_index].release()
        print "WRONG!!! Not locked index = %d" % (sock_index)
        
    socks[sock_index].close()
    
    ticks = 1
    while True:
        try:
            tcpsock = socket(AF_INET, SOCK_STREAM) 
            tcpsock.connect(('8.8.8.8', 53)) 
            break
        except Exception, e:
            tcpsock.close()
            print "Failed to connect, try %d seconds later..." % (ticks)
            time.sleep(ticks)
            ticks = ticks * 2
    
    socks[sock_index] = tcpsock
    
    gl_sock_locks[sock_index].release()
    
    print "sock #" + str(sock_index) + " renewed!"

class algo2_get_sock():
    def __init__(self, socks):
        self.socks = socks
        self.max = len(socks)
        
    def get_sock(self, index):
        
        for i in range(self.max):
            
            # Try lock this sock. Should fail if already locked due to renewing or in-use
            if ( gl_sock_locks[index].acquire(False)):
            
#                gl_sock_locks[index].release()
#                print "Get sock #%d, and lock it for use" % (index)
                return index
            else:
                print "SOck #%d already locked by others" % (index)
                
            if ( index == self.max - 1):
                index = 0
            else:
                index = index + 1
        
        # No available sock for transmitting
        return None
            
class algo_get_sock():
    def __init__(self, socks):
        self.socks = socks
        self.max = len(socks)
        
    def get_sock(self, index):
        
        valid_sock = False
        for i in range(self.max):
            if ( index == self.max - 1):
                index = 0
            else:
                index = index + 1
            
            # Try lock this sock. Should fail if it is locked due to renewing
            if ( gl_sock_locks[index].acquire(False)):
            
                gl_sock_locks[index].release()
                valid_sock = True
                break
        
        if ( valid_sock):
            print "Get sock #" + str(index)
            return index
        else:
            # No available sock for transmitting
            return None
    
class TCPConnectionHandler():
    def __init__(self, socknum):
        self.socks = []
        for i in range(socknum):
            tcpsock = socket(AF_INET, SOCK_STREAM)
            tcpsock.connect(('8.8.8.8', 53))
            self.socks.append(tcpsock)
        self.index = 0
        self.max = socknum
        
        for i in range(self.max):
            gl_sock_locks.append( threading.Lock())
            
        self.get_sock_handler = algo2_get_sock(self.socks)
            
    # Manage socket status of connection
    def run(self):
        pass
    
    def get_sock(self):
        # Access self.index by mutex
        with gl_mutex:
            new_index = self.get_sock_handler.get_sock(self.index) 
            
            if new_index != None:
                self.index = new_index
                return self.socks[self.index]
            else:
                # No available sock for transmitting
                print "NO available socks"
                return None
        
    def cleanup_sock(self, sock, release):
        if ( sock == None):
            return
            
        if ( release == False):
            return
            
        with gl_mutex:
            for i in range(self.max):
                if (self.socks[i] == sock):
                    if ( release == False):
                        pass
                    else:
                        # Need to release existing lock
                        gl_sock_locks[i].release()
                        print "In cleanup sock #%d released" % (i)
                        
                    break
                
    
    def renew_sock(self, sock):
        # Assume sock already in locked state, so do not bother to release it
        self.cleanup_sock( sock, False)
        
        # Find sock in pool, close and create a new one
        gl_mutex.acquire()
        
        for i in range(self.max):
            if (self.socks[i] == sock):
                print "Prepare to renew sock#%d" % (i)
#                gl_sock_locks[i].acquire()
                sock.close()
                
                # Kick off a new thread for sock renewal
                thread.start_new_thread(thread_get_sock, (self.socks, i))
                
                break
                
        gl_mutex.release()
    
glTCPConnectionHandler = TCPConnectionHandler(8)
glCache = {}

class LocalDNSHandler(BaseRequestHandler):
    def setup(self):
        remote_server = '8.8.8.8'
        self.dnsserver = (remote_server, DEF_REMOTE_PORT)

    def handle(self):
        print "New request received from " + str(self.client_address)
        data, socket = self.request

        if ( data.find("TEST") ==0 ):
            print data
            socket.sendto("HELLO", self.client_address)
            return
            
#        print "Request !!!!!!!!!!"
#        print repr(data)
        domain = self.getRequestDomain( data)
        if ( domain in glCache.keys()) :
            # Need to replace the response with the request ID
            rspdata = self.getRequestID(data) + glCache[domain][2:]
        else: 
            rspdata = self._getResponseTCP(data)
            
            if ( rspdata == None):
                return
                
            if ( len(rspdata) == 0):
                return
            glCache[domain] = rspdata
            
#        print repr(rspdata)
        socket.sendto(rspdata, self.client_address)
            
    def getRequestID( self, data):
        return data[0:2]
            
    def getRequestDomain( self, data):
#        print "Length of request packet is %d" % (len(data))
        if ( len(data) > 13):
            question = data[12:]
            domain = self.parseDomain( question)
            domain = domain[0:len(domain)-1]
            return domain
        else:
            return None
    
    def parseDomain( self, data):
        result = ""
        length = len(data)
#        print repr(data)
        partLength = unpack('H', '\x00'+data[0])[0]
#        print "Length = %d" % ( partLength)
        
        if ( partLength == 0):
            return result
        
        remained = data[partLength +1 :]
        result = data[1:partLength+1] + "."
        result += self.parseDomain(remained)
        
        return result 
         
            
    def _getResponseTCP(self, data):
    
        global glTCPConnectionHandler
        
        udpResp = None
        
        tcpData = pack('!h', len(data)) + data
                            
        for i in range(3):
            tcpsock = glTCPConnectionHandler.get_sock()
            try:
                tcpsock.setblocking(1)
                tcpsock.sendall(tcpData)
                
                tcpsock.settimeout(1)
                resp = tcpsock.recv(512)
                
                print str(time.time()) + ": Response from TCP, len = " + str(len(resp))
                if ( len(resp) == 0):
                    glTCPConnectionHandler.renew_sock( tcpsock)
                else:
                    # release the lock on tcpsock
                    glTCPConnectionHandler.cleanup_sock(tcpsock, True)
                    udpResp = resp[2:]
                    break
                
            except Exception, e:
                print str(time.time()) + ": Failed !!!!!!!!!!!!!!!!! %d times" % (i+1)
                glTCPConnectionHandler.renew_sock( tcpsock)
        
        return udpResp
        
class LocalDNSServer(ThreadingUDPServer):
    pass

def main():
    global gl_remote_server
    gl_remote_server = '8.8.8.8'
    
    dnsserver = LocalDNSServer((DEF_LOCAL_HOST, DEF_LOCAL_PORT), LocalDNSHandler)
    dnsserver.serve_forever()

if __name__ == '__main__':
    daemonize()
    main()

